import argparse
import jsonlines

def pubtabnet_split(origin_jsonl_path, train_jsonl_path, val_jsonl_path):
    with jsonlines.open(origin_jsonl_path, "r") as f:
        with jsonlines.open(train_jsonl_path, "w") as train_f:
            for data in f:
                if data['split'] == "train":
                    train_f.write(data)

    with jsonlines.open(origin_jsonl_path, "r") as f:
        with jsonlines.open(val_jsonl_path, "w") as val_f:
            for data in f:
                if data['split'] == "val":
                    val_f.write(data)

if __name__ == "__main__":
    # 路径建议均使用绝对路径
    parser = argparse.ArgumentParser()
    parser.add_argument('--origin_jsonl_path', type=str, help="Path to origin jsonl file", default=None)
    parser.add_argument('--train_jsonl_path', type=str, help="Path to train data jsonl file", default=None)
    parser.add_argument('--val_jsonl_path', type=str, help="Path to eval data jsonl file", default=None)
    args = parser.parse_args()

    pubtabnet_split(args.origin_jsonl_path, args.train_jsonl_path, args.val_jsonl_path)